from casevo.agent_base import AgentBase
import numpy as np
import pandas as pd
from typing import Dict, Any, List

class OpinionAgent(AgentBase):
    """观点智能体类，用于模拟用户在社交网络中的观点表达行为"""
    
    def __init__(self, unique_id: int, model: Any, description: str, context: Dict[str, Any], 
                 user_features: pd.DataFrame = None, sentiment_data: pd.DataFrame = None):
        """
        初始化观点智能体
        
        Args:
            unique_id: 智能体唯一标识
            model: 模型环境
            description: 智能体描述
            context: 上下文信息
            user_features: 用户特征数据
            sentiment_data: 情感分析数据
        """
        super().__init__(unique_id, model, description, context)
        
        # 从处理后的数据初始化智能体属性
        if user_features is not None and unique_id in user_features['user_id'].values:
            user_data = user_features[user_features['user_id'] == unique_id].iloc[0]
            self.opinion = user_data['opinion_tendency']
            self.confidence = user_data['emotional_stability']
            self.activity_score = user_data['activity_score']
            self.influence_score = user_data['influence_score']
            self.group_id = user_data['group_id']
        else:
            # 如果没有数据，使用随机值
            self.opinion = np.random.uniform(-1.0, 1.0)
            self.confidence = np.random.uniform(0.5, 1.0)
            self.activity_score = np.random.uniform(0.3, 1.0)
            self.influence_score = np.random.uniform(0.3, 1.0)
            self.group_id = np.random.randint(1, 4)
        
        # 初始化其他属性
        self.social_pressure = 0.0
        self.is_silent = False
        self.memory = []
        
        # 设置思考链
        self.setup_chain({
            'opinion_update': {
                'steps': [
                    'observe_neighbors',
                    'calculate_pressure',
                    'update_opinion',
                    'decide_speak'
                ]
            }
        })
        
    def step(self) -> None:
        """执行智能体的一个时间步"""
        # 获取邻居
        neighbors = self.model.grid.get_neighbors(self.unique_id)
        
        # 观察邻居观点
        neighbor_opinions = [agent.opinion for agent in neighbors]
        
        # 计算社交压力（考虑影响力）
        if neighbor_opinions:
            weighted_opinions = []
            for agent in neighbors:
                weight = agent.influence_score * (1 if agent.group_id == self.group_id else 0.5)
                weighted_opinions.append(agent.opinion * weight)
            avg_opinion = np.mean(weighted_opinions)
            self.social_pressure = self.confidence * 0.3 * (avg_opinion - self.opinion)
        
        # 更新观点（考虑活跃度）
        self.opinion += self.social_pressure * self.activity_score
        self.opinion = np.clip(self.opinion, -1.0, 1.0)
        
        # 决定是否表达观点（考虑自信度）
        if abs(self.opinion) < self.confidence * 0.5:
            self.is_silent = True
        else:
            self.is_silent = False
            
        # 记录到记忆
        self.memory.append({
            'opinion': self.opinion,
            'pressure': self.social_pressure,
            'is_silent': self.is_silent
        })
        
        # 保持记忆大小
        if len(self.memory) > 10:
            self.memory.pop(0)
            
    def get_state(self) -> Dict[str, Any]:
        """
        获取智能体当前状态
        
        Returns:
            Dict: 包含智能体状态的字典
        """
        return {
            'agent_id': self.unique_id,
            'opinion': self.opinion,
            'confidence': self.confidence,
            'activity_score': self.activity_score,
            'influence_score': self.influence_score,
            'group_id': self.group_id,
            'social_pressure': self.social_pressure,
            'is_silent': self.is_silent,
            'memory': self.memory.copy()
        } 